import torch
from torch.nn import Softmax
from torch.nn import CrossEntropyLoss, MSELoss
from typing import Optional, Tuple

from transformers import RobertaForSequenceClassification

from transformers.modeling_outputs import SequenceClassifierOutput

from dataclasses import dataclass

@dataclass
class SequenceClassifierOutputWithLastLayer(SequenceClassifierOutput):
    """
    用于存储序列分类器的输出，包括损失、 logits、最后一个隐藏状态、隐藏状态序列和注意力分数。

    参数:
    - loss: 可选的Tensor，表示分类损失。
    - logits: Tensor，表示分类器的输出logits。
    - last_hidden_state: Tensor，表示模型最后一个层的隐藏状态。
    - hidden_states: 可选的元组，包含模型所有隐藏层的状态。
    - attentions: 可选的元组，包含模型所有层的注意力分数。
    """
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    last_hidden_state: torch.FloatTensor = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


class RobertaForContrastiveClassification(RobertaForSequenceClassification):
    """
    Roberta模型的派生类，用于对比学习分类任务。

    该类在RobertaForSequenceClassification的基础上添加了softmax层，用于输出经过softmax激活的logits,计算LCE。

    参数:
    - config: 模型配置对象，定义了模型的结构和参数。
    """
    _keys_to_ignore_on_load_missing = [r"position_ids"]

    def __init__(self, config):
        super().__init__(config)
        self.soft_max = Softmax(dim=1)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        """
        前向传播方法，用于执行模型的计算。

        参数:
        - input_ids: 可选的Tensor，表示输入序列的token id。
        - attention_mask: 可选的Tensor，表示输入序列的注意力mask。
        - token_type_ids: 可选的Tensor，表示输入序列的token类型id。
        - position_ids: 可选的Tensor，表示输入序列的位置id。
        - head_mask: 可选的Tensor，表示注意力头的mask。
        - inputs_embeds: 可选的Tensor，表示输入序列的嵌入向量。
        - labels: 可选的Tensor，表示输入序列的标签。
        - output_attentions: 可选的布尔值，控制是否输出注意力分数。
        - output_hidden_states: 可选的布尔值，控制是否输出隐藏层状态。
        - return_dict: 可选的布尔值，控制是否以SequenceClassifierOutputWithLastLayer对象的形式返回结果。

        返回:
        - 如果return_dict为True，则返回SequenceClassifierOutputWithLastLayer对象；
        - 如果return_dict为False，则返回一个元组，包含softmax logits和其他额外输出。
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = outputs[0]
        # [cls]转化为logits
        logits = self.classifier(sequence_output)

        loss = None
        if labels is not None:
            if self.num_labels == 1:
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                # 计算LCE
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        softmax_logits = self.soft_max(logits)

        if not return_dict:
            output = (softmax_logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutputWithLastLayer(
            loss=loss,
            logits=softmax_logits,
            last_hidden_state=sequence_output,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
